import gym_pih
import gymnasium as gym
from seals import classic_control
import time
import numpy as np
import time
import sys
import cv2
import math
from gym_pih.hardware.calculate_slope import calculate_slope
from gym_pih.wrapper.record_success_rate import RecordSuccessRateWrapper
from matplotlib import pyplot as plt

env=gym.make('serl/pih-v2')
env=RecordSuccessRateWrapper(env)
# cv2.namedWindow('image',cv2.WINDOW_NORMAL)
try:
    action=np.zeros(3)
    z_force=[]
    observation,_=env.reset()
    for i in range(200):
        start_time=time.time()
        action[:2]=np.random.uniform(-1,1,2)
        action[2]=np.random.uniform(-1,0)
        # action=np.zeros(3)
        observation,reward,terminated,truncated,info=env.step(action)
        # print(reward)
        z_force.append(env.ft.forceTorque[2])
        sys.stdout.write("\r"+ " " * 100 + "\r")
        # sys.stdout.write(str(observation["state"]["force"]*10))
        sys.stdout.write("control frequency:"+str(1.0/(time.time()-start_time)))
        sys.stdout.flush()
        _,image=env.Realsense_cam.GetImage()
        _,side_image=env.Imagingsource_cam.GetImage()
        slope_image=side_image[130:,250:1020]

        _,slope,binary_cleaned=calculate_slope(slope_image,is_return_process_image=True)
    

        # print(result_deg, "°")
        # 显示图像

        side_image=side_image[40:720, 360:1000]
        side_image=cv2.resize(side_image,(64,64))

        cv2.namedWindow("side_image",cv2.WINDOW_NORMAL)
        cv2.resizeWindow("side_image", 320, 320)
        cv2.imshow("side_image",side_image)
        

        cv2.namedWindow("Fitted Line", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("Fitted Line", 1280, 720)
        cv2.imshow("Fitted Line", binary_cleaned)
        

        image_clip=image[150:214,210:274]
        image_clip=cv2.resize(image_clip,(64,64))
        cv2.namedWindow('image_clip',cv2.WINDOW_NORMAL)
        cv2.imshow('image_clip',image_clip)
        cv2.waitKey(1)

        if terminated or truncated:
            observation,_=env.reset()

    print('done')
    plt.plot(z_force)
    plt.show()
finally:
    env.reset()
    env.close()
    cv2.destroyAllWindows()